﻿using Fst.SemiRing;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Fst.Operations
{
    /// <summary>
    /// Remove epsilon operation.
    /// 
    /// @author John Salatas <jsalatas@users.sourceforge.net>
    /// </summary>
    public class RmEpsilon
    {
        /**
        /// Default Constructor
         */
        private RmEpsilon() {
        }

        /**
        /// Put a new state in the epsilon closure
         */
        private static void put(State fromState, State toState, float weight,
                Dictionary<State, float>[] cl) 
        {
            Dictionary<State, float> tmp = cl[fromState.getId()];
            if (tmp == null) {
                tmp = new Dictionary<State, float>();
                cl[fromState.getId()] = tmp;
            }
            tmp.Add(toState, weight);
        }

        /**
        /// Add a state in the epsilon closure
         */
        private static void add(State fromState, State toState, float weight,
                Dictionary<State, float>[] cl, Semiring semiring) 
        {
            float old = getPathWeight(fromState, toState, cl);
            if (old == float.NaN) {
                put(fromState, toState, weight, cl);
            } else {
                put(fromState, toState, semiring.plus(weight, old), cl);
            }

        }

        /**
        /// Calculate the epsilon closure
         */
        private static void calcClosure(Fst fst, State state,
                Dictionary<State, float>[] cl, Semiring semiring) 
        {
            State s = state;

            float pathWeight;
            int numArcs = s.getNumArcs();
            for (int j = 0; j < numArcs; j++) {
                Arc a = s.getArc(j);
                if ((a.getIlabel() == 0) && (a.getOlabel() == 0)) {
                    if (cl[a.getNextState().getId()] == null) {
                        calcClosure(fst, a.getNextState(), cl, semiring);
                    }
                    if (cl[a.getNextState().getId()] != null) 
                    {
                        foreach (State pathFinalState in cl[a.getNextState().getId()].Keys) 
                        {
                            pathWeight = semiring.times(
                                    getPathWeight(a.getNextState(), pathFinalState,
                                            cl), a.getWeight());
                            add(state, pathFinalState, pathWeight, cl, semiring);
                        }
                    }
                    add(state, a.getNextState(), a.getWeight(), cl, semiring);
                }
            }
        }

        /**
        /// Get an epsilon path's cost in epsilon closure
         */
        public static float getPathWeight(State _in, State _out,
                Dictionary<State, float>[] cl) 
        {
            if (cl[_in.getId()] != null) 
            {
                return cl[_in.getId()][_out];
            }

            return float.NaN;
        }

        /**
        /// Removes epsilon transitions from an fst.
        /// 
        /// It return a new epsilon-free fst and does not modify the original fst
        /// 
        /// @param fst the fst to remove epsilon transitions from
        /// @return the epsilon-free fst
         */
        public static Fst get(Fst fst) {
            if (fst == null) {
                return null;
            }

            if (fst.getSemiring() == null) {
                return null;
            }

            Semiring semiring = fst.getSemiring();

            Fst res = new Fst(semiring);

            Dictionary<State, float>[] cl = new Dictionary<State, float>[fst.getNumStates()];
            State[] oldToNewStateMap = new State[fst.getNumStates()];
            State[] newToOldStateMap = new State[fst.getNumStates()];

            int numStates = fst.getNumStates();
            for (int i = 0; i < numStates; i++) {
                State s = fst.getState(i);
                // Add non-epsilon arcs
                State newState = new State(s.getFinalWeight());
                res.addState(newState);
                oldToNewStateMap[s.getId()] = newState;
                newToOldStateMap[newState.getId()] = s;
                if (newState.getId() == fst.getStart().getId()) {
                    res.setStart(newState);
                }
            }

            for (int i = 0; i < numStates; i++) {
                State s = fst.getState(i);
                // Add non-epsilon arcs
                State newState = oldToNewStateMap[s.getId()];
                int numArcs = s.getNumArcs();
                for (int j = 0; j < numArcs; j++) {
                    Arc a = s.getArc(j);
                    if ((a.getIlabel() != 0) || (a.getOlabel() != 0)) {
                        newState.addArc(new Arc(a.getIlabel(), a.getOlabel(), a
                                .getWeight(), oldToNewStateMap[a.getNextState()
                                .getId()]));
                    }
                }

                // Compute e-Closure
                if (cl[s.getId()] == null) {
                    calcClosure(fst, s, cl, semiring);
                }
            }

            // augment fst with arcs generated from epsilon moves.
            numStates = res.getNumStates();
            for (int i = 0; i < numStates; i++) {
                State s = res.getState(i);
                State oldState = newToOldStateMap[s.getId()];
                if (cl[oldState.getId()] != null) {
                    foreach (State pathFinalState in cl[oldState.getId()].Keys) {
                        State s1 = pathFinalState;
                        if (s1.getFinalWeight() != semiring.zero()) {
                            s.setFinalWeight(semiring.plus(s.getFinalWeight(),
                                    semiring.times(getPathWeight(oldState, s1, cl),
                                            s1.getFinalWeight())));
                        }
                        int numArcs = s1.getNumArcs();
                        for (int j = 0; j < numArcs; j++) {
                            Arc a = s1.getArc(j);
                            if ((a.getIlabel() != 0) || (a.getOlabel() != 0)) {
                                Arc newArc = new Arc(a.getIlabel(), a.getOlabel(),
                                        semiring.times(a.getWeight(),
                                                getPathWeight(oldState, s1, cl)),
                                        oldToNewStateMap[a.getNextState().getId()]);
                                s.addArc(newArc);
                            }
                        }
                    }
                }
            }

            res.setIsyms(fst.getIsyms());
            res.setOsyms(fst.getOsyms());

            Connect.apply(res);

            return res;
        }
    }
}
